import os, sys
import argparse
import json
import cv2
from PIL import Image
from pylab import *

sys.path.append("../")

from utils.utils import get_yolo_boxes, makedirs
from utils.bbox import draw_boxes, get_boxes
from keras.models import load_model
# from tqdm import tqdm
import numpy as np
import time

# add by suyongsheng
import tensorflow as tf
from multiprocessing import cpu_count


class Yolo:
    def __init__(self, config_path):
        tf.Session(config=tf.ConfigProto(device_count={"CPU": cpu_count()}, inter_op_parallelism_threads=1,
                                         intra_op_parallelism_threads=2))
        with open(config_path) as config_buffer:
            self.config = json.load(config_buffer)

        # self.net_h, self.net_w = 416, 416
        self.net_h, self.net_w = 288, 288
        self.obj_thresh, self.nms_thresh = 0.5, 0.45
        os.environ['CUDA_VISIBLE_DEVICES'] = self.config['train']['gpus']
        self.infer_model = load_model(self.config['train']['saved_weights_name'])

    def predict(self, image, from_file=False, show_im=False):
        if from_file:
            image = cv2.imread(image)
        time_start = time.time()
        # predict the bounding boxes
        boxes = \
            get_yolo_boxes(self.infer_model, [image], self.net_h, self.net_w, self.config['model']['anchors'],
                           self.obj_thresh, self.nms_thresh)[0]
        time_end = time.time()
        print('yolo predict cost time=', time_end - time_start)
        xyc_pack = get_boxes(image, boxes, self.config['model']['labels'], self.obj_thresh)
        if show_im:
            image = draw_boxes(image, boxes, self.config['model']['labels'], self.obj_thresh)
            imshow(image)
            show()
        return xyc_pack

    def predict_dir(self, dir_in, dir_out):
        file_list=[]
        for dir in os.listdir(dir_in):
            for file in os.listdir(os.path.join(dir_in,dir)):
                file_full_name = os.path.join(dir_in,dir,file)
                if os.path.splitext(file_full_name)[1] in ['.jpg', '.png', 'JPEG']:
                    file_list.append(file_full_name)
        i = 0
        if os.path.exists(dir_out) is False:
            os.makedirs(dir_out)
        for file in file_list:
            image = cv2.imread(file)
            boxes = \
            get_yolo_boxes(self.infer_model, [image], self.net_h, self.net_w, self.config['model']['anchors'],
                           self.obj_thresh, self.nms_thresh)[0]
            sys.stdout.write('\r>> processing im %d/%d ' % (
                i , len(file_list)))
            i += 1
            sys.stdout.flush()


            draw_boxes(image, boxes, self.config['model']['labels'], self.obj_thresh)
            cv2.imwrite(dir_out+"/" + file.split('/')[-1], np.uint8(image))


def case1():
    model_yolo = Yolo("./config_eval.json")
    image = cv2.imread("../0.jpg")
    model_yolo.predict(image, show_im=False)
    model_yolo.predict(image, show_im=False)
    model_yolo.predict(image, show_im=False)
    model_yolo.predict(image, show_im=False)
    model_yolo.predict(image, show_im=True)

def case2(args):
    model_yolo = Yolo("./config_eval.json")
    model_yolo.predict_dir("/home/leo/Downloads/compare/origin/",
                           "/home/leo/Downloads/compare/origin_train/")

def case3(args):
    model_yolo = Yolo(args.conf)
    model_yolo.predict_dir(args.input,
                           args.output)


if __name__ == '__main__':
    argparser = argparse.ArgumentParser(description='train and evaluate YOLO_v3 model on any dataset')
    argparser.add_argument('-c', '--conf', help='path to configuration file')
    argparser.add_argument('-i', '--input', help='path to configuration file')
    argparser.add_argument('-o', '--output', help='path to configuration file')

    args = argparser.parse_args()
    case3(args)
    # case1()